"""
数据配置文件
"""

import os
from pathlib import Path

class DataConfig:
    """数据配置类"""
    
    # 基础路径配置
    BASE_DIR = Path(__file__).parent.parent
    DATA_DIR = BASE_DIR / "data"
    RAW_DIR = DATA_DIR / "raw"
    PROCESSED_DIR = DATA_DIR / "processed"
    SPLITS_DIR = DATA_DIR / "splits"
    
    # 原始数据文件
    EXCEL_FILE = RAW_DIR / "TCIA-CPTAC-CCRCC_v11_20230818-nbia-digest.xlsx"
    IMAGES_DIR = RAW_DIR / "images"
    
    # 数据分割比例
    TRAIN_RATIO = 0.7
    VALIDATION_RATIO = 0.2
    TEST_RATIO = 0.1
    
    # 图像处理配置
    IMAGE_SIZE = (224, 224)  # 统一图像尺寸
    CHANNELS = 3  # 图像通道数
    NORMALIZE_MEAN = [0.485, 0.456, 0.406]  # ImageNet标准化均值
    NORMALIZE_STD = [0.229, 0.224, 0.225]   # ImageNet标准化标准差
    
    # 数据增强配置
    AUGMENTATION_CONFIG = {
        "train": {
            "horizontal_flip": True,
            "vertical_flip": False,
            "rotation": 15,
            "brightness": 0.1,
            "contrast": 0.1,
            "saturation": 0.1,
            "hue": 0.1,
            "scale": (0.8, 1.2),
            "translate": (0.1, 0.1),
            "normalize": True
        },
        "val": {
            "normalize": True
        },
        "test": {
            "normalize": True
        }
    }
    
    # 类别配置
    CLASS_NAMES = [
        "normal",           # 正常
        "cyst",            # 囊肿
        "tumor",           # 肿瘤
        "stone",           # 结石
        "infection"        # 感染
    ]
    
    CLASS_MAPPING = {
        "normal": 0,
        "cyst": 1,
        "tumor": 2,
        "stone": 3,
        "infection": 4
    }
    
    # 数据加载配置
    BATCH_SIZE = 16
    NUM_WORKERS = 4
    PIN_MEMORY = True
    
    # 缓存配置
    CACHE_DIR = DATA_DIR / "cache"
    USE_CACHE = True
    
    # 数据验证配置
    VALIDATION_CONFIG = {
        "check_image_integrity": True,
        "check_label_consistency": True,
        "min_image_size": (100, 100),
        "max_file_size_mb": 50
    }
    
    @classmethod
    def create_directories(cls):
        """创建必要的目录"""
        directories = [
            cls.DATA_DIR,
            cls.RAW_DIR,
            cls.PROCESSED_DIR,
            cls.SPLITS_DIR,
            cls.IMAGES_DIR,
            cls.CACHE_DIR
        ]
        
        for directory in directories:
            directory.mkdir(parents=True, exist_ok=True)
    
    @classmethod
    def get_image_paths(cls, split="train"):
        """获取指定分割的图像路径"""
        split_dir = cls.SPLITS_DIR / split
        if not split_dir.exists():
            raise FileNotFoundError(f"分割目录不存在: {split_dir}")
        
        image_paths = []
        for ext in ['*.jpg', '*.jpeg', '*.png', '*.bmp', '*.tiff']:
            image_paths.extend(split_dir.rglob(ext))
        
        return sorted(image_paths)
    
    @classmethod
    def get_label_file(cls, split="train"):
        """获取标签文件路径"""
        return cls.SPLITS_DIR / f"{split}_labels.csv"
    
    @classmethod
    def get_metadata_file(cls):
        """获取元数据文件路径"""
        return cls.PROCESSED_DIR / "metadata.csv"

